#!/usr/bin/env python3
"""
测试外围背景替换增强功能

该脚本用于验证 PeripheralBackgroundReplacement 类的功能，
展示如何使用该增强方法模拟复杂背景。
"""

import os
import sys
import cv2
import numpy as np
import argparse
from pathlib import Path

# 导入外围背景替换增强类
from datasets import PeripheralBackgroundReplacement


def create_sample_edge_map(image_shape, edge_thickness=3):
    """创建一个简单的边缘图用于测试
    
    Args:
        image_shape: (height, width) 图像尺寸
        edge_thickness: 边缘线条粗细
    
    Returns:
        edge_map: (H, W) 边缘图, uint8, 值域 0-255
    """
    h, w = image_shape
    edge_map = np.zeros((h, w), dtype=np.uint8)
    
    # 绘制矩形边缘
    margin = h // 4
    cv2.rectangle(edge_map, (margin, margin), (w - margin, h - margin), 255, edge_thickness)
    
    # 绘制对角线
    cv2.line(edge_map, (margin, margin), (w - margin, h - margin), 255, edge_thickness)
    cv2.line(edge_map, (w - margin, margin), (margin, h - margin), 255, edge_thickness)
    
    return edge_map


def visualize_augmentation(image, edge_gt, augmented_image, save_path=None):
    """可视化增强结果
    
    Args:
        image: 原始图像
        edge_gt: 边缘标注
        augmented_image: 增强后的图像
        save_path: 保存路径，如果为None则不保存
    """
    # 准备可视化图像
    # 1. 原始图像
    img_display = image.copy()
    if len(img_display.shape) == 2:
        img_display = cv2.cvtColor(img_display, cv2.COLOR_GRAY2BGR)
    elif img_display.shape[-1] == 1:
        img_display = cv2.cvtColor(img_display.squeeze(), cv2.COLOR_GRAY2BGR)
    
    # 2. 边缘图（彩色显示）
    edge_display = cv2.cvtColor(edge_gt, cv2.COLOR_GRAY2BGR)
    
    # 3. 增强后的图像
    aug_display = augmented_image.copy()
    if len(aug_display.shape) == 2:
        aug_display = cv2.cvtColor(aug_display, cv2.COLOR_GRAY2BGR)
    elif aug_display.shape[-1] == 1:
        aug_display = cv2.cvtColor(aug_display.squeeze(), cv2.COLOR_GRAY2BGR)
    
    # 4. 边缘叠加在原图上
    edge_overlay = img_display.copy()
    edge_overlay[edge_gt > 127] = [0, 255, 0]  # 绿色边缘
    
    # 5. 边缘叠加在增强图上
    aug_overlay = aug_display.copy()
    aug_overlay[edge_gt > 127] = [0, 255, 0]  # 绿色边缘
    
    # 组合成一行
    h, w = img_display.shape[:2]
    
    # 添加标题
    def add_title(img, title):
        img_with_title = np.zeros((h + 40, w, 3), dtype=np.uint8)
        img_with_title[40:, :] = img
        cv2.putText(img_with_title, title, (10, 25), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        return img_with_title
    
    img1 = add_title(img_display, "Original Image")
    img2 = add_title(edge_display, "Edge Ground Truth")
    img3 = add_title(edge_overlay, "Edge Overlay")
    img4 = add_title(aug_display, "Augmented Image")
    img5 = add_title(aug_overlay, "Aug + Edge Overlay")
    
    # 水平拼接
    result = np.hstack([img1, img2, img3, img4, img5])
    
    if save_path:
        cv2.imwrite(save_path, result)
        print(f"可视化结果已保存到: {save_path}")
    
    return result


def test_peripheral_bg_augmentation(args):
    """测试外围背景替换增强功能"""
    
    print("=" * 80)
    print("外围背景替换增强测试")
    print("=" * 80)
    
    # 1. 加载测试图像
    if not os.path.isfile(args.test_image):
        print(f"错误: 测试图像不存在: {args.test_image}")
        return
    
    image = cv2.imread(args.test_image, cv2.IMREAD_COLOR)
    if image is None:
        print(f"错误: 无法读取图像: {args.test_image}")
        return
    
    print(f"\n1. 加载测试图像: {args.test_image}")
    print(f"   - 图像尺寸: {image.shape}")
    
    # 2. 创建或加载边缘图
    if args.edge_map and os.path.isfile(args.edge_map):
        edge_gt = cv2.imread(args.edge_map, cv2.IMREAD_GRAYSCALE)
        print(f"2. 加载边缘图: {args.edge_map}")
    else:
        edge_gt = create_sample_edge_map(image.shape[:2], edge_thickness=args.edge_thickness)
        print(f"2. 创建示例边缘图")
    
    print(f"   - 边缘图尺寸: {edge_gt.shape}")
    
    # 3. 初始化外围背景替换增强器
    if not os.path.isdir(args.background_dir):
        print(f"\n错误: 背景图像目录不存在: {args.background_dir}")
        return
    
    print(f"\n3. 初始化外围背景替换增强器")
    print(f"   - 背景图像目录: {args.background_dir}")
    print(f"   - 应用概率: {args.prob}")
    print(f"   - 边缘保护范围: {args.edge_dilation}px")
    print(f"   - 外围平滑范围: {args.peripheral_erosion}px")
    
    try:
        augmentor = PeripheralBackgroundReplacement(
            background_dir=args.background_dir,
            prob=args.prob,
            edge_dilation=args.edge_dilation,
            peripheral_erosion=args.peripheral_erosion
        )
    except Exception as e:
        print(f"\n错误: 初始化增强器失败: {e}")
        return
    
    # 4. 应用增强
    print(f"\n4. 应用外围背景替换增强")
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 生成多个样本以展示随机性
    num_samples = args.num_samples
    for i in range(num_samples):
        print(f"   - 生成样本 {i+1}/{num_samples}...")
        
        augmented_image, _ = augmentor(image.copy(), edge_gt.copy())
        
        # 可视化结果
        save_path = os.path.join(args.output_dir, f"peripheral_bg_aug_sample_{i+1}.png")
        visualize_augmentation(image, edge_gt, augmented_image, save_path)
    
    print(f"\n5. 测试完成！")
    print(f"   - 输出目录: {args.output_dir}")
    print(f"   - 生成样本数: {num_samples}")
    print("\n" + "=" * 80)


def main():
    parser = argparse.ArgumentParser(description='测试外围背景替换增强功能')
    
    parser.add_argument('--test_image',
                        type=str,
                        required=True,
                        help='测试图像路径')
    parser.add_argument('--edge_map',
                        type=str,
                        default=None,
                        help='边缘图路径（如果不提供则自动生成示例边缘图）')
    parser.add_argument('--background_dir',
                        type=str,
                        required=True,
                        help='背景图像目录路径')
    parser.add_argument('--output_dir',
                        type=str,
                        default='peripheral_bg_test_results',
                        help='输出目录路径')
    parser.add_argument('--num_samples',
                        type=int,
                        default=5,
                        help='生成的增强样本数量')
    parser.add_argument('--prob',
                        type=float,
                        default=1.0,
                        help='应用概率（测试时建议设为1.0以确保每次都应用）')
    parser.add_argument('--edge_dilation',
                        type=int,
                        default=20,
                        help='边缘保护范围（像素）')
    parser.add_argument('--peripheral_erosion',
                        type=int,
                        default=15,
                        help='外围区域平滑范围（像素）')
    parser.add_argument('--edge_thickness',
                        type=int,
                        default=3,
                        help='示例边缘图的线条粗细（仅在不提供edge_map时使用）')
    
    args = parser.parse_args()
    
    test_peripheral_bg_augmentation(args)


if __name__ == '__main__':
    main()
